from cProfile import label
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import matplotlib.ticker as ticker

class PlotData(object):
    def __init__(self, x:list, y:list, fig_name:str='training error', legend:list[str]=['train loss', 'valid loss', 'valid acc']) -> None:
        self.x = x
        self.y = y
        self.fig_name = fig_name
        self.fig, self.ax = plt.subplots()
        self.ax.set_title(self.fig_name, fontsize=16, fontweight='bold')
        self.ax.set_xlabel('Epoch', fontsize=12)
        self.right = self.ax.twinx() # 设置右坐标
        self.ax.set_ylabel('loss', fontsize=24)
        self.right.set_ylabel('accuracy', fontsize=24)
        self.legend  = legend
        self.ax.grid(True, alpha=0.7)
        self.right.grid(True, alpha=0.7)
        self.ax.set_xbound(min(self.x)-1, max(self.x)+1)
        y_max = 0
        for y in self.y:
            y_max = max(y) if max(y) > y_max else y_max
        self.ax.set_ybound(0, int(y_max)+1)
        self.ax.yaxis.set_major_locator(ticker.MultipleLocator(10))
        self.ax.yaxis.set_minor_locator(ticker.MultipleLocator(2))
        self.ax.yaxis.set_minor_formatter(ticker.ScalarFormatter())
        self.right.yaxis.set_major_locator(ticker.MultipleLocator(10))
        self.right.yaxis.set_minor_locator(ticker.MultipleLocator(2))

    def __call__(self) -> None:
        for k, y in enumerate(self.y[0:-2]):
            self.ax.plot(self.x, y, label=self.legend[k])
        # plt.tight_layout(pad=23)
        self.right.plot(self.x, self.y[-1], label=self.legend[-1], color='C4')
        self.right.plot(self.x, self.y[-2], label=self.legend[-2], color='C9')
        self.ax.legend(loc='upper center')
        self.right.legend(loc='upper right')
        plt.show(block=True)


if __name__ == "__main__":
    plot = PlotData([i for i in range(10)], [[i for i in range(10)], 
                                             [i for i in range(3, 13)]])
    plot()
    print(matplotlib.__version__)